/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Declaration of the stellarmodel class.
/// \ingroup sfrhist

// $Id$


#ifndef __model__
#define  __model__

#include <vector>
#include <string> 

#include "blitz/array.h"
#include "particle_set.h"
#include "boost/function.hpp"
#include "mcrx-units.h"
#include "vecops.h"
#include "sphgrid.h"

namespace mcrx {
  class star_particle_ptr;
  class stellarmodel;
  class simple_stellarmodel;
  class mappings_stellarmodel;
  class mappings_sb99_model;
  class nbody_data_grid;
}

class Preferences;
class sfrhist;
namespace CCfits {
  class FITS;
};


/** Extension of the particle_ptr class that also includes a reference
    to the SED array for the particle. */
class mcrx::star_particle_ptr : 
  public mcrx::particle_ptr<mcrx::star_particle_set> {
  array_2* sed_;
public:
  star_particle_ptr(star_particle_set& s, int n, array_2& ss) :
    particle_ptr<star_particle_set>(s,n), sed_(&ss) {};
  array_2& sed() const {return *sed_;};
};


/** Abstract base class for stellar models, that are able to generate
    SEDs as a fnction of age, metallicity, and other parameters. */
class mcrx::stellarmodel {
public:
  virtual ~stellarmodel() {};
  virtual array_1 get_SED(star_particle_set&, int, T_float) const =0;
  /** Calculate the SEDs of all the particles contained in the vector. */
  virtual void calculate_SEDs(const std::vector<star_particle_ptr>&, 
			      T_float) const =0;

  /** Returns the current fraction of mass still present in stars for
      a population of a certain time and metallicity.  */
  virtual double get_mass_fraction (T_float time, T_float metallicity) const=0;
  virtual const array_1 delta_lambda() const=0;
  virtual const array_1& lambda() const=0;
  virtual int n_lambda() const=0;

  virtual void set_pressure_functor(boost::function<std::pair<T_float,std::pair<T_float, T_float> >(const vec3d&)> f) {};

  /*** Resamples the SED in the stellar model onto the wavelength
       vector supplied. Normally this just involves interpolating, but
       the class is free to do something more advanced like actually
       integrate to make sure energy is preserved. */
  virtual void resample(const array_1& lambda)=0;

  virtual const T_unit_map& units () const=0;

  virtual void write_fits_parameters (const std::string&) const=0; 
};

/** Contains the spectral energy distributions from a stellar
    population synthesis model. The spectrum of an instantaneous burst
    as a function of time is used for calculating the spectral energy
    distributions of the particles by convolving with the
    star-formation rate history. \ingroup sfrhist */
class mcrx::simple_stellarmodel : public mcrx::stellarmodel {
private:
  Preferences& prefs;
  std::string file_name; ///< The stellar model filename.
  
  /// The array holding the SEDs.  Time is first index, metallicity
  /// second, and lambda third.
  array_3 sed; 
  array_2 current_mass; ///< The current fraction of mass remaining in stars
  array_2 L_bol; ///< The bolometric luminosity in the SED array.
  std::vector<T_float>  times; ///< The times of the entries in the SED array.
  std::vector<T_float>  middle_times; ///< The time midway between SED entries.
  /// The metallicities of the entries in the SED array.
  std::vector<T_float>  metallicities; 
  /// The metallicities midway between SED entries.
  std::vector<T_float>  middle_metallicities;
  std::vector<T_float> lambdas; ///< The wavelengths in the SED array.
  /** Array keeping the wavelength values. This just references the
      data in axes[l_axis], so must be kept up to date if that is
      changed. */
  array_1 alambda_;
  // Counters to track how many particles have age or metallicity outside the
  // range covered by the templates.
  mutable int n_gt_max_age;
  mutable int n_lt_min_age;
  mutable int n_gt_max_met;
  mutable int n_lt_min_met;
  
  // some useful parameters
  bool log_flux; ///< If true, the SED array contains log10 of the luminosity.
  /// The initial stellar mass for which the SED is given.
  float reference_mass_;

  T_unit_map units_;

  void calculate_L_bol();

public:
  simple_stellarmodel(Preferences& p);

  /** Copies the STELLARMODEL and LAMBDA HDU's to the output file. */
  virtual void write_fits_parameters (const std::string&) const; 

  /** The main access function, which returns the SED associated with
      a certain time and metallicity.  */
  virtual array_1 get_SED(star_particle_set& s, int index, T_float time) const;
  array_1 get_SED (T_float age, T_float metallicity) const;

  /** Calculate the SEDs of all the particles contained in the vector. */
  virtual void calculate_SEDs(const std::vector<star_particle_ptr>&,
			      T_float) const;

  virtual void resample(const array_1& lambda);

  /** Returns the current fraction of mass still present in stars for
      a population of a certain time and metallicity.  */
  double get_mass_fraction (T_float time, T_float metallicity) const;
  int n_times () const {return sed.extent(blitz::firstDim );};
  int n_metallicities () const {return sed.extent(blitz::secondDim );};
  int n_lambda () const {return sed.extent(blitz::thirdDim );};
  
  /** Returns the \f$\Delta\lambda\f$ vector indicating the width of
      the "wavelength bins" in the SED. This is handy when integrating
      the SED over wavelength. */
  virtual const array_1 delta_lambda () const {
    return delta_quantity(alambda_);};
  virtual const array_1& lambda () const {return alambda_;};

  /// Unit retrieval.
  virtual const T_unit_map& units () const {return units_;};
};


/** Contains the spectral energy distributions from a MAPPINGS
    photoionization HII model.  SEDs for the HII regions (ie young
    stellar particles) depend on the ISM pressure in addition to age
    and metalicity. \ingroup sfrhist */
class mcrx::mappings_stellarmodel : public mcrx::stellarmodel {
private:
  Preferences& prefs;

  /// This functor is used to get the ISM pressure at a specified position.
  boost::function<std::pair<T_float,std::pair<T_float, T_float> >(const vec3d&)> pressure_functor;

  /// Fixed value of M_cl used.
  T_float Mcl_;
  /// Fixed value of S used.
  T_float S_;
  /// PDR fraction used
  T_float pdr_f_;
  /// Determines whether we ask for SEDs with a fixed Mcl or S
  bool fixed_mcl_; 
  /// Determines whether effective temp is used to calculate P/k.
  bool use_teff_;
  /// Determines whether overlapping MAPPINGS particles should get the
  /// aggregate cluster mass
  bool do_clustering_;
  /// Determines whether MAPPINGS particles have their radius set to r_s.
  bool set_particle_radius_;

  /// The number of dimensions of the mappings SED.
  static const int nvar_=5;
  static const int s_axis=0;
  static const int p_axis=1;
  static const int z_axis=2;
  static const int t_axis=3;
  static const int l_axis=4;
  static const int c_axis=5; // only for the temp array

  // Counters to track how many particles have age or metallicity outside the
  // range covered by the templates.
  mutable long n_gt_max_age;
  mutable long n_lt_min_age;
  mutable long n_gt_max_met;
  mutable long n_lt_min_met;
  mutable long n_gt_max_S;
  mutable long n_lt_min_S;
  mutable long n_gt_max_P;
  mutable long n_lt_min_P;

  /** The array holding the MAPPINGS SEDs.  The axes are given by the
      x_axis constants.  */
  blitz::Array<T_float, nvar_> sed; 

  /// The axis values in the different dimensions.
  std::vector<std::vector<T_float> > axes; 
  /// The midpoints between axis values.
  std::vector<std::vector<T_float> > middle_axes; 
  /** Array keeping the wavelength values. This just references the
      data in axes[l_axis], so must be kept up to date if that is
      changed. */
  array_1 alambda_;
  
  T_unit_map units_;

  // keeps the radii of the pdr
  blitz::Array<T_float, 2> r_pdr_;

  /** gcc doesn't emit debugging info for constructors so we have to
      do stupid shit like this... Called by the constructor. */
  void stupid ();
  array_2 read_brent_file(const std::string& fn);
  array_2 read_brent_file2(const std::string& fn);

  /** Calculates P/k for the specified grid cell. */
  T_float calculate_pressure(const nbody_data_grid::T_cell* const c) const;

  /** Returns the PDR radius for the model with specified Pk and S. */
  T_float get_pdr_radius(T_float Pk, T_float Sparam) const;

  /** Calculates the mass in the PDR and compares it to the mass
      available around the particle, and prints the quantities.
      Returns r_s, the radius enclosing the PDR mass. */
  T_float check_pdr_mass(T_float rho_gas, T_float z_gas,
			 T_float Pk, T_float age, 
			 T_float radius, T_float S, 
			 T_float mass) const;

public:
  mappings_stellarmodel(Preferences&);

  /** Sets the functor used for returning the grid cell containing the
      particle so we can calculate the ISM pressure. */
  void set_pressure_functor(boost::function<std::pair<T_float,std::pair<T_float, T_float> >(const vec3d&)> f) {
    pressure_functor = f;};


  /** This does nothing because mappings information is in the
      preferences keywords that are already written. We should perhaps
      write some MAPPINGS-specific info? */
  void write_fits_parameters (const std::string& output_file_name) const {};

  /** Calculate the SEDs of all the particles contained in the
      vector. This includes doing a clustering analysis of the
      particles to determine if a larger cluster mass should be
      used. */
  virtual void calculate_SEDs(const std::vector<star_particle_ptr>&,
			      T_float) const;

  /** The main access function, which returns the SED associated with
      a certain time and metallicity.  */
  virtual array_1 get_SED(star_particle_set& s, int index, T_float time) const;
  array_1 get_SED(T_float age, T_float metallicity, 
		  T_float Pk, T_float Sparam) const;

  double get_mass_fraction (T_float time, T_float metallicity) const {
    return 1;};

  virtual void resample(const array_1& lambda) {assert(0);};

  int n_pressures () const {return axes[p_axis].size();};
  int n_times () const {return axes[t_axis].size();};
  T_float oldest () const {return axes[t_axis].back();};
  int n_metallicities () const {return axes[z_axis].size();};
  int n_lambda () const {return axes[l_axis].size();};
  
  /* Returns the wavelength interval over which each of the points in
      the SED are valid.  This is handy when integrating over
      wavelength. */
  virtual const array_1 delta_lambda () const {
    return delta_quantity(alambda_);};
  virtual const array_1& lambda () const {return alambda_;};

  /// Unit retrieval.
  virtual const T_unit_map& units () const {return units_;};
};


/** A combined MAPPINGS and Starburst99 model that uses MAPPINGS SEDs
    for the ages available, otherwise the SB99 SEDs. */
class mcrx::mappings_sb99_model : public stellarmodel {
private:
  simple_stellarmodel sb99;
  mappings_stellarmodel mappings;

  /// If false (set with keyword use_mappings_seds), only sb99 seds are used
  bool use_mappings;

public:
  mappings_sb99_model(Preferences&);

  void set_pressure_functor(boost::function<std::pair<T_float,std::pair<T_float, T_float> >(const vec3d&)> f) {
    mappings.set_pressure_functor(f);};

  /** Calculate the SEDs of all the particles contained in the
      vector. This splits the particles by time and sends them to the
      mappings or sb99 models. */
  virtual void calculate_SEDs(const std::vector<star_particle_ptr>&,
			      T_float) const;

  /** This function asks either of the models for the SED. */
  virtual array_1 get_SED(star_particle_set&, int, T_float) const;
  array_1 get_SED(T_float age, T_float metallicity, 
		  T_float Pk, T_float Sparam) const;

  virtual void resample(const array_1& lambda) {assert(0);};

  virtual double get_mass_fraction (T_float time, T_float metallicity) const {
    return sb99.get_mass_fraction(time, metallicity);};
  virtual const array_1 delta_lambda() const { return sb99.delta_lambda();};
  virtual const array_1& lambda() const { return sb99.lambda();};
  virtual int n_lambda() const { return sb99.n_lambda();};

  virtual void write_fits_parameters (const std::string& s) const {
    sb99.write_fits_parameters(s);
    mappings.write_fits_parameters(s);};

  /// Unit retrieval.
  virtual const T_unit_map& units () const {return sb99.units();};
};

#endif
